# 引入相关的库
import random
import math
import pandas as pd
import numpy as np

# 模型类
class LFM(object):
    def __init__(self, F,file_path,test,alpha=0.01, lmbd=0.01, max_iter=300):
        self.F = F
        self.P = dict()  # R=PQ^T，代码中的Q相当于博客中Q的转置
        self.Q = dict()
        self.alpha = alpha
        self.lmbd = lmbd
        self.max_iter = max_iter
        self.file_path = file_path
        self.rating_data = self.get_rating_data()
        self.mu = 0.0
        self.bu = dict()
        self.bi = dict()

        self.test = test
        
        '''随机初始化矩阵P和Q'''
        cnt = 0
        for user, rates in self.rating_data.items():
            self.P[user] = [random.random() / math.sqrt(self.F) for x in range(self.F)]
            self.bu[user]=0
            cnt += len(rates)
            for item, rate in rates.items():
                self.mu += rate
                if item not in self.Q:
                    self.Q[item] = [random.random() / math.sqrt(self.F) for x in range(self.F)]
                self.bi[item] = 0
        
        self.mu /= cnt
        
        
    def get_rating_data(self):
        rating_data = dict()
        # 训练数据集
        with open(self.file_path, "r") as f:
            for line in f.readlines()[1:]:
                user_id, item_id, score,_ = line.split(",")
                if float(score) >= 1.0:
                    score = 8.5
                else:
                    score = 1.5
                #print(float(score))
                rating_data.setdefault(user_id,dict())
                rating_data[user_id][item_id] = score
        print("训练数据加载完毕")
        return rating_data


    """
    随机梯度下降法训练参数P和Q
    :return: 
    """
    def train(self):
        for step in range(self.max_iter):
            loss = 0.0
            for user, rates in self.rating_data.items():
                for item, rui in rates.items():
                    hat_rui = self.predict(user, item)
                    err_ui = rui - hat_rui
                    loss += abs(err_ui)
                    # 更新偏置
                    self.bu[user] += self.alpha * (err_ui - self.lmbd * self.bu[user])
                    self.bi[item] += self.alpha * (err_ui - self.lmbd * self.bi[item])
                    for f in range(self.F):
                        self.P[user][f] += self.alpha * (err_ui * self.Q[item][f] - self.lmbd * self.P[user][f])
                        self.Q[item][f] += self.alpha * (err_ui * self.P[user][f] - self.lmbd * self.Q[item][f])
            self.alpha *= 0.9
            
            print("{}轮:loss：{}".format(step,loss))
            if (step+1)%10==0:
                r0 = []
                for user, item in test:
                    try:
                        rate = lfm0.predict(user, item)
                    except:
                        rate = 7.5

                    if rate >= 5.0:
                        r0.append(1)
                    else:
                        r0.append(0)

                save_csv(r0,"new_{}_{}".format(step+1,loss))

    def predict(self, user, item):
        return sum(self.P[user][f] * self.Q[item][f] for f in range(self.F))+ self.bu[user] + self.bi[item] + self.mu



def save_csv(cnt, name):
    ids = [i for i in range(1, len(cnt) + 1)]
    df_data = []
    for i in range(len(ids)):
        temp = []
        temp.append(ids[i])
        temp.append(cnt[i])
        df_data.append(temp)

    df = pd.DataFrame(np.array(df_data), columns=['id', 'rating'])
    df.to_csv("temp/{}.csv".format(name), index=False)
    print("{} 写入完成".format(name))
    
def get_test_data(pt):
    test = list()
    with open(pt,"r") as f:
        for line in f.readlines()[1:]:
            user_id, item_id = line[:-1].split(",")
            test.append((user_id, item_id))
    test = np.array(test)
    return test
    
if __name__ =="__main__":
    p0 = "dataSets/train.csv"
    pt = "dataSets/test.csv"

    test = get_test_data(pt)

    lfm0 = LFM(F=60,file_path=p0,test=test,max_iter=310)

    # 模型训练
    lfm0.train()

    r0 = []
    for user, item in test:
        try:
            rate = lfm0.predict(user, item)
        except:
            rate = 7.5

        if rate >= 5.0:
            r0.append(1)
        else:
            r0.append(0)

    save_csv(r0, "pianzhi_lfm_2")
    
